import tensorflow as tf
import time
import cv2
import sys
import os
import shutil
import numpy as np
slim = tf.contrib.slim
from PIL import Image

# tfrecord的位置
PATH_RES = [
    r'./data_sour/pj_vehicle_train_00000-of-00004.tfrecord',
    r'./data_sour/pj_vehicle_train_00001-of-00004.tfrecord',
    r'./data_sour/pj_vehicle_train_00002-of-00004.tfrecord',
    r'./data_sour/pj_vehicle_train_00003-of-00004.tfrecord',
    r'./data_sour/pj_vehicle_validation_00000-of-00004.tfrecord',
    r'./data_sour/pj_vehicle_validation_00001-of-00004.tfrecord',
    r'./data_sour/pj_vehicle_validation_00002-of-00004.tfrecord',
    r'./data_sour/pj_vehicle_validation_00003-of-00004.tfrecord',
    ]
# 图片存放位置
PATH_DES = [
    r'./data/images',
    r'./data/images',
    r'./data/images',
    r'./data/images',
    r'./data/images',
    r'./data/images',
    r'./data/images',
    r'./data/images',
    ]

PATH = list(zip(PATH_RES, PATH_DES))


def tfrecord2jpg(path_res, path_des):
    print('tfrecords_files to be transformed:', path_res)
    reader = tf.TFRecordReader()
    start_time = int(time.time())
    prev_time = start_time
    idx = 0

    filename_queue = tf.train.string_input_producer([path_res], num_epochs=1)

    # 从 TFRecord 读取内容并保存到 serialized_example 中
    _, serialized_example = reader.read(filename_queue)

    # 读取 serialized_example 的格式
    features = tf.parse_single_example(
        serialized_example,
        features={
            'image/encoded': tf.FixedLenFeature([], tf.string),
            'image/format': tf.FixedLenFeature([], tf.string),
            'image/class/label': tf.FixedLenFeature([], tf.int64),
            'image/height': tf.FixedLenFeature([], tf.int64),
            'image/width': tf.FixedLenFeature([], tf.int64),

        })
    slim.tfexample_decoder.Image(image_key='image/encoded',
                                 format_key='image/format'),

    # 使用 tf.image.decode_jpeg对jpg格式图像进行解码，对应tf.gfile读取图像，method_1
    image_target = tf.image.decode_jpeg(features['image/encoded'])

    # 使用tf.decode_raw将字符串解析成图像对应的像素数组，对应Image.open读取图像，method_2
    # image_target = tf.decode_raw(features['image/encoded'], tf.uint8)
    # images = tf.decode_raw(features['image/encoded'], tf.uint8)

    format = tf.decode_raw(features['image/format'], tf.uint8)
    labels = tf.cast(features['image/class/label'], tf.int64)
    height = tf.cast(features['image/height'], tf.int64)
    width = tf.cast(features['image/width'], tf.int64)

    print('Extracting {} has just started.'.format(path_res))
    with tf.Session() as sess:
        # 启动多线程
        sess.run(tf.local_variables_initializer())
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)
        while not coord.should_stop():
            try:
                outlabel, outformat, outimg, outwidth, outheight = sess.run([labels, format, image_target,
                                                                             width, height])
                print("---->outimg.shape:", outimg.shape, "---label:", outlabel, "--->format:", outformat, "--->width:",
                      outwidth, "--->height:", outheight)
            except tf.errors.OutOfRangeError:
                print("Turn to next folder.")
                break

            image = Image.fromarray(outimg)

            name = str(idx)+"_"+str(outlabel) + "_" + str(outwidth) + "_" + str(outheight)
            idx = idx + 1
            image.save(path_des+"/reverse_%s.jpg" % name)  # 保存部分图片查看
            current_time = int(time.time())
            lasting_time = current_time - start_time
            interval_time = current_time - prev_time
            if interval_time >= 0.1:
                sys.stdout.flush()
                sys.stdout.write("\rGenerating the {}-th image: {},\
                                    lasting {} seconds".format(
                                    idx,
                                    path_des +
                                    str(idx) + '_' +
                                    str(outlabel) + "_"+ str(outwidth)+"_"+ str(outheight)+'.jpg',
                                    lasting_time))
                prev_time = current_time
        coord.request_stop()
        coord.join(threads)

def main():
    # get empty directory
    for i in range(len(PATH)):
        # if os.path.isdir(PATH_DES[i]):
        #     if os.listdir(PATH_DES[i]):
        #         print("---->PATH_DES[i]",PATH_DES[i])
        #         shutil.rmtree(PATH_DES[i])
        #         os.mkdir(PATH_DES[i])
        # else:
        #     print(PATH_DES[i])
        #     os.mkdir(PATH_DES[i])
        tfrecord2jpg(PATH_RES[i], PATH_DES[i])


if __name__ == "__main__":
    main()
